import lzma
import os
import struct
from pathlib import Path
from io import BufferedReader, BytesIO

# --- Constants ---
HEADER_LEN = 1024
OWN_KEY_LEN = 17
PUB_KEY_LEN = 1170494464 # Not directly used for allocation size like in Rust
PUB_KEY_LEN_MAGNIFICATION = 16
MAGIC_HEADER = bytes([
    0x7c, 0xd5, 0x32, 0xeb, 0x86, 0x02, 0x7f, 0x4b, 0xa8, 0xaf, 0xa6, 0x8e, 0x0f, 0xff, 0x99,
    0x14, 0x00, 0x04, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
])
PUB_KEY_MEND = bytes([
    0xB8, 0xD5, 0x3D, 0xB2, 0xE9, 0xAF, 0x78, 0x8C, 0x83, 0x33, 0x71, 0x51, 0x76, 0xA0,
    0xCD, 0x37, 0x2F, 0x3E, 0x35, 0x8D, 0xA9, 0xBE, 0x98, 0xB7, 0xE7, 0x8C, 0x22, 0xCE,
    0x5A, 0x61, 0xDF, 0x68, 0x69, 0x89, 0xFE, 0xA5, 0xB6, 0xDE, 0xA9, 0x77, 0xFC, 0xC8,
    0xBD, 0xBD, 0xE5, 0x6D, 0x3E, 0x5A, 0x36, 0xEF, 0x69, 0x4E, 0xBE, 0xE1, 0xE9, 0x66,
    0x1C, 0xF3, 0xD9, 0x02, 0xB6, 0xF2, 0x12, 0x9B, 0x44, 0xD0, 0x6F, 0xB9, 0x35, 0x89,
    0xB6, 0x46, 0x6D, 0x73, 0x82, 0x06, 0x69, 0xC1, 0xED, 0xD7, 0x85, 0xC2, 0x30, 0xDF,
    0xA2, 0x62, 0xBE, 0x79, 0x2D, 0x62, 0x62, 0x3D, 0x0D, 0x7E, 0xBE, 0x48, 0x89, 0x23,
    0x02, 0xA0, 0xE4, 0xD5, 0x75, 0x51, 0x32, 0x02, 0x53, 0xFD, 0x16, 0x3A, 0x21, 0x3B,
    0x16, 0x0F, 0xC3, 0xB2, 0xBB, 0xB3, 0xE2, 0xBA, 0x3A, 0x3D, 0x13, 0xEC, 0xF6, 0x01,
    0x45, 0x84, 0xA5, 0x70, 0x0F, 0x93, 0x49, 0x0C, 0x64, 0xCD, 0x31, 0xD5, 0xCC, 0x4C,
    0x07, 0x01, 0x9E, 0x00, 0x1A, 0x23, 0x90, 0xBF, 0x88, 0x1E, 0x3B, 0xAB, 0xA6, 0x3E,
    0xC4, 0x73, 0x47, 0x10, 0x7E, 0x3B, 0x5E, 0xBC, 0xE3, 0x00, 0x84, 0xFF, 0x09, 0xD4,
    0xE0, 0x89, 0x0F, 0x5B, 0x58, 0x70, 0x4F, 0xFB, 0x65, 0xD8, 0x5C, 0x53, 0x1B, 0xD3,
    0xC8, 0xC6, 0xBF, 0xEF, 0x98, 0xB0, 0x50, 0x4F, 0x0F, 0xEA, 0xE5, 0x83, 0x58, 0x8C,
    0x28, 0x2C, 0x84, 0x67, 0xCD, 0xD0, 0x9E, 0x47, 0xDB, 0x27, 0x50, 0xCA, 0xF4, 0x63,
    0x63, 0xE8, 0x97, 0x7F, 0x1B, 0x4B, 0x0C, 0xC2, 0xC1, 0x21, 0x4C, 0xCC, 0x58, 0xF5,
    0x94, 0x52, 0xA3, 0xF3, 0xD3, 0xE0, 0x68, 0xF4, 0x00, 0x23, 0xF3, 0x5E, 0x0A, 0x7B,
    0x93, 0xDD, 0xAB, 0x12, 0xB2, 0x13, 0xE8, 0x84, 0xD7, 0xA7, 0x9F, 0x0F, 0x32, 0x4C,
    0x55, 0x1D, 0x04, 0x36, 0x52, 0xDC, 0x03, 0xF3, 0xF9, 0x4E, 0x42, 0xE9, 0x3D, 0x61,
    0xEF, 0x7C, 0xB6, 0xB3, 0x93, 0x50,
])

# --- Lazy Loading Public Key ---
_pub_key_data = None

def _get_asset_path(filename="kugou_key.xz"):
    """Gets the absolute path to the asset file."""
    # Assumes assets directory is in the same directory as this script
    script_dir = Path(__file__).parent.resolve()
    return script_dir / "assets" / filename

def _load_pub_key():
    """Loads and decompresses the public key from the xz file."""
    global _pub_key_data
    if _pub_key_data is None:
        key_path = _get_asset_path()
        if not key_path.exists():
             raise FileNotFoundError(f"KuGou key file not found at {key_path}. "
                                     "Make sure 'kugou_key.xz' is in the 'assets' directory.")
        try:
            with lzma.open(key_path, "rb") as f:
                # Read the entire decompressed key. The size calculation from Rust
                # suggests the expected decompressed size.
                # expected_size = PUB_KEY_LEN // PUB_KEY_LEN_MAGNIFICATION
                # _pub_key_data = f.read(expected_size)
                # Reading all is safer if the size constant isn't perfectly accurate
                _pub_key_data = f.read()
                # Optional: Verify size if needed
                # expected_size = PUB_KEY_LEN // PUB_KEY_LEN_MAGNIFICATION
                # if len(_pub_key_data) != expected_size:
                #     print(f"Warning: Decoded key size {len(_pub_key_data)} differs from expected {expected_size}")

        except Exception as e:
            raise RuntimeError(f"Failed to decode the KuGou key from {key_path}: {e}")
    return _pub_key_data

def _get_pub_key_segment(start_pos, length):
    """
    Retrieves the relevant segment of the public key.
    Note: This differs slightly from Rust's direct indexing. We fetch the
    required range based on the start position and length of the data chunk.
    """
    pub_key = _load_pub_key()
    start_index = start_pos // PUB_KEY_LEN_MAGNIFICATION
    # Calculate end index needed based on the last byte's position
    end_pos = start_pos + length
    end_index = end_pos // PUB_KEY_LEN_MAGNIFICATION

    # Ensure indices are within bounds
    start_index = max(0, start_index)
    end_index = min(len(pub_key) -1 , end_index) # -1 because index is inclusive

    # Return the segment slice. Add 1 to end_index because Python slicing excludes the end.
    return pub_key[start_index : end_index + 1]


# --- Decoder Class ---
class KuGouDecoder:
    """
    Decoder for Kugou KGM/VPR files. Behaves like a readable file object.
    """
    def __init__(self, file_obj: BufferedReader):
        """
        Initializes the decoder.

        Args:
            file_obj (BufferedReader): The binary read stream of the original encrypted file.
        """
        self.origin = file_obj
        self.own_key = bytearray(OWN_KEY_LEN)
        self.pos = 0 # Position in the *decrypted* stream

        # Read and validate header
        header = self.origin.read(HEADER_LEN)
        if len(header) != HEADER_LEN or not header.startswith(MAGIC_HEADER):
            # Reset position if header check fails, so the file can be used elsewhere
            try:
                self.origin.seek(0)
            except Exception:
                pass # Seek might fail on non-seekable streams
            raise ValueError("Invalid KGM data: Header mismatch or too short.")

        # Extract own_key (offset 0x1c, length 16 bytes)
        self.own_key[:16] = header[0x1c : 0x1c + 16]
        # The 17th byte is often 0, matching the Rust initialization.
        # If VPR requires a different 17th byte, logic needs adjustment here.

        # Ensure the public key is loaded (or load it now)
        _load_pub_key()

    def read(self, size=-1):
        """
        Reads and decrypts data from the underlying file stream.

        Args:
            size (int): The number of bytes to read. Reads till EOF if -1 or omitted.

        Returns:
            bytes: The decrypted data, or b'' if EOF is reached.
        """
        # Read encrypted data from the original file
        # Note: We don't use self.origin.tell() because the decryption logic
        # relies on the absolute position from the start of the *encrypted* data,
        # which we track with self.pos.
        encrypted_data = self.origin.read(size)
        if not encrypted_data:
            return b'' # EOF

        data_len = len(encrypted_data)
        decrypted_data = bytearray(data_len)
        pub_key_segment = _get_pub_key_segment(self.pos, data_len)
        pub_key_start_index = self.pos // PUB_KEY_LEN_MAGNIFICATION

        for i in range(data_len):
            current_pos = self.pos + i
            byte = encrypted_data[i]

            # Apply own_key XOR
            own_key_byte = self.own_key[current_pos % OWN_KEY_LEN]
            med8 = own_key_byte ^ byte
            med8 ^= (med8 & 0x0f) << 4

            # Apply pub_key XOR
            pub_key_index_in_segment = (current_pos // PUB_KEY_LEN_MAGNIFICATION) - pub_key_start_index
            # Boundary check for safety, though _get_pub_key_segment should provide enough
            if pub_key_index_in_segment < 0 or pub_key_index_in_segment >= len(pub_key_segment):
                 # This case should ideally not happen if segment logic is correct
                 # print(f"Warning: pub_key_index_in_segment out of bounds: {pub_key_index_in_segment}")
                 # Fallback or error? Using 0 for now, might indicate a problem.
                 pub_key_byte = 0
            else:
                pub_key_byte = pub_key_segment[pub_key_index_in_segment]

            mend_byte = PUB_KEY_MEND[current_pos % len(PUB_KEY_MEND)]
            msk8 = mend_byte ^ pub_key_byte
            msk8 ^= (msk8 & 0xf) << 4

            decrypted_data[i] = med8 ^ msk8

            # VPR mask diff logic was commented out in Rust, omitted here too.
            # if raw_ext == "vpr":
            #     decrypted_data[i] ^= VprMaskDiff[current_pos % 17]

        self.pos += data_len
        return bytes(decrypted_data)

    def close(self):
        """Closes the underlying file object."""
        if self.origin:
            self.origin.close()
            self.origin = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False # Don't suppress exceptions

    # --- Methods matching Rust Decoder trait (for conceptual mapping) ---
    @staticmethod
    def decodeable_length_interval():
        """ conceptually similar to Rust, returns min/max size for a valid file """
        # This isn't strictly needed for the Python logic but maps to the Rust trait
        # A file must be at least HEADER_LEN to be possibly valid. Max size isn't checked here.
        return (HEADER_LEN, float('inf'))

    @classmethod
    def try_new(cls, file_obj: BufferedReader):
        """
        Attempts to create a decoder. Returns None if the header is invalid.
        Similar to try_new in Rust.
        """
        try:
            # Remember current position
            current_pos = file_obj.tell()
            decoder = cls(file_obj)
            # If successful, seek back to the position *after* the header
            # because the __init__ consumed it.
            file_obj.seek(current_pos + HEADER_LEN)
            return decoder
        except (ValueError, struct.error, EOFError):
            # Reset position if initialization failed
            try:
                file_obj.seek(current_pos)
            except Exception:
                 pass # Seek might fail
            return None
        except FileNotFoundError as e: # Catch key file not found specifically
            print(f"Error: {e}")
            return None


# --- Factory function (like src/decoder/mod.rs) ---
def new_decoder(file_obj: BufferedReader):
    """
    Attempts to create a KuGou decoder for the given file object.
    Returns the decoder instance or None if the format is not recognized.
    """
    # Currently only supports KuGou
    return KuGouDecoder.try_new(file_obj)

# --- Test function (similar to Rust test) ---
def _test_decode():
    # Assumes test files are in 'src/assets' relative to the original project root
    # Adjust path as needed for your test setup
    test_file_path = Path("../../src/assets/test_kugou_kgm.dat") # Adjust relative path
    right_file_path = Path("../../src/assets/test_kugou_kgm_right.dat") # Adjust relative path

    if not test_file_path.exists() or not right_file_path.exists():
        print("Skipping test: Test files not found.")
        print(f"Expected KGM at: {test_file_path.resolve()}")
        print(f"Expected DAT at: {right_file_path.resolve()}")
        return

    print(f"Testing with: {test_file_path}")

    decoded_data = BytesIO()
    try:
        with open(test_file_path, "rb") as f_in:
            # Use the factory function
            decoder = new_decoder(f_in)
            if not decoder:
                 print("Test failed: Could not create decoder.")
                 return

            with decoder: # Use context manager
                while True:
                    chunk = decoder.read(16 * 1024) # Read in chunks
                    if not chunk:
                        break
                    decoded_data.write(chunk)
    except Exception as e:
        print(f"Test failed during decoding: {e}")
        import traceback
        traceback.print_exc()
        return

    try:
        with open(right_file_path, "rb") as f_right:
            right_dat = f_right.read()
    except Exception as e:
        print(f"Test failed: Could not read reference file {right_file_path}: {e}")
        return

    audio = decoded_data.getvalue()
    print(f"Decoded size: {len(audio)}, Reference size: {len(right_dat)}")

    if audio == right_dat:
        print("Test PASSED!")
    else:
        print("Test FAILED: Decoded data does not match reference data.")
        # Optional: Write decoded data for inspection
        # with open("decoded_test_output.bin", "wb") as f_out:
        #     f_out.write(audio)
        # print("Wrote decoded output to decoded_test_output.bin")


if __name__ == '__main__':
    # Run the test function if the script is executed directly
    _test_decode() 